perf: Extend WindowTopN to dense_rank - #23869
Conversation
github-actions
Bot
added
optimizer
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23869 +/- ##
==========================================
- Coverage 80.72% 80.70% -0.02%
==========================================
Files 1090 1095 +5
Lines 370384 373167 +2783
Branches 370384 373167 +2783
==========================================
+ Hits 298975 301163 +2188
- Misses 53590 54002 +412
- Partials 17819 18002 +183 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
comphead
changed the title
dense_rank
|
@kumarUjjawal @kosiew can you please review this PR? |
kumarUjjawal
left a comment
kumarUjjawal
left a comment
There was a problem hiding this comment.
Hi @SubhamSinghal I have left some comments please let me know what you think.
| match udf.fun().name() { | ||
| "row_number" => Some(WindowFnKind::RowNumber), | ||
| "rank" => Some(WindowFnKind::Rank), | ||
| "dense_rank" => Some(WindowFnKind::DenseRank), |
There was a problem hiding this comment.
Could we skip this rewrite when another expression in the same window operator needs following rows? The input is pruned before all window expressions run, so a sibling lead() returns the wrong value at the retained boundary. We can add a dense_rank with lead regression and only rewrite when every sibling expression is safe under tail pruning.
There was a problem hiding this comment.
Fixed in 4530bba. The rewrite now bails unless every window expr is ROW_NUMBER/RANK/DENSE_RANK over the same PARTITION BY/ORDER BY. LEAD and aggregates are excluded. Added EXPLAIN + correctness SLT regressions for SUM and LEAD siblings.
| // For RANK / DENSE_RANK: an empty ORDER BY makes every row tie | ||
| // at rank 1 — the optimization is degenerate (we'd retain the | ||
| // entire input) and tie storage would be unbounded. | ||
| if matches!(fn_kind, WindowFnKind::Rank | WindowFnKind::DenseRank) |
There was a problem hiding this comment.
This guard misses orderings that become empty after partition keys are removed. DENSE_RANK() OVER (PARTITION BY pk ORDER BY pk) reaches PartitionedTopKExec without order expressions and panics. Please validate the effective order keys and add a regression.
There was a problem hiding this comment.
Fixed in 4530bba. Replaced order_by().is_empty() with a structural guard: sort_exec.expr().len() <= partition_prefix_len → bail. Catches PARTITION BY pk ORDER BY pk (and no-ORDER-BY ROW_NUMBER). Added SLT regression.
| entries.push(TieEntry { | ||
| batch: batch.clone(), | ||
| row_indices: run_indices, | ||
| batch_bytes: input_batch_bytes, |
There was a problem hiding this comment.
records the full source-batch size for every retained group even though the batch clones share Arrow buffers. With many partitions, memory reservation is heavily inflated and can fail with ResourcesExhausted. This expands #23326 to the normal dense-rank path; please account each source batch once or compact the retained rows.
There was a problem hiding this comment.
that will leaves the normal dense-rank path counting one shared batch once per retained group across every partition. This can greatly inflate the reservation and cause false ResourcesExhausted errors. Could we fix the shared-batch accounting here, or keep this rewrite disabled until #23326 is fixed?
| .groups | ||
| .keys() | ||
| .map(|key| key.as_slice()) | ||
| .max() |
There was a problem hiding this comment.
Finding the maximum scans all K keys for every new distinct value. On mostly distinct input this makes the rewrite O(N × K) and can be much slower for large K than the sort it replaces.
There was a problem hiding this comment.
I had kept it unoptimised thinking K will be small. Fixed in 4530bba. Cache the admission boundary (max_key); the O(K) scan runs only when stale (after eviction/fill), so the common "worse than boundary → drop" path is O(1).
| if is_smaller { | ||
| // Evict the entire max-key group and drop the cached | ||
| // max (recomputed on the next Case C). | ||
| let evicted_key = state.max_key.take().expect("max key present"); |
There was a problem hiding this comment.
The cache only makes rejected candidates O(1). Every accepted candidate invalidates it, so input arriving from worst to best still scans all K keys for each new value and remains O(N × K). The benchmarks use K=2 and do not exercise this case. Could we use an ordered structure or add a large-K improving-order benchmark?
| entries.push(TieEntry { | ||
| batch: batch.clone(), | ||
| row_indices: run_indices, | ||
| batch_bytes: input_batch_bytes, |
There was a problem hiding this comment.
that will leaves the normal dense-rank path counting one shared batch once per retained group across every partition. This can greatly inflate the reservation and cause false ResourcesExhausted errors. Could we fix the shared-batch accounting here, or keep this rewrite disabled until #23326 is fixed?
|
@2010YOUY01 Can you take a look at this whenever you get the chance. |
3 participants
Which issue does this PR close?
ROW_NUMBER < 5/ TopK #6899. Completes the ROW_NUMBER + RANK + DENSE_RANK trio requested by the issue (ROW_NUMBER shipped in Perf: Window topn optimisation #21479; RANK is in flight as PR A).Rationale for this change
Naive plan for
Filter(dr <= K) → BoundedWindowAggExec(DENSE_RANK) → SortExec(pk, ob) → inputsorts the full input even though only rows at the K distinct-smallest ob values per partition are needed. DENSE_RANK is stable under tail-pruning, so the rewrite preserves ranks for every retained row.What changes are included in this PR?
Are these changes tested?
Yes
h2o bench — Q24–Q29 added; top-2 on 10M-row
large, 3-iter avg vsenable_window_topn=false:Are there any user-facing changes?
No breaking changes. Rule fires on DENSE_RANK when
datafusion.optimizer.enable_window_topn = true; flag staysfalseby default